package cardsdk

import (
	"fmt"
	"math/big"

	"gitee.com/fhe-search/nfhedecd/NFHE_Function"
	"gitee.com/fhe-search/nfhedecd/interfaceAPI"
)

func FHEUint16Enc(uint16ClearText uint16) ([]byte, error) {
	return NFHE_Function.NFHE_ENC(uint(uint16ClearText))
}

func FHEUint16Dec(uint16EncText []byte) (uint16, error) {
	if len(uint16EncText) != 128 {
		return 0, fmt.Errorf("EncText length isn't 128")
	}
	// NFHE_Function.NFHE_DEC会改变入参，所以这里需要拷贝一份
	uint16EncText = append([]byte{}, uint16EncText...)
	decText, err := NFHE_Function.NFHE_DEC(uint16EncText)
	return uint16(decText), err
}

func FHEUtf8StrEnc(utf8ClearText string) ([]byte, error) {
	uint16ClearTexts := []uint16{}
	for i := 0; i < len(utf8ClearText)/2; i++ {
		b1 := utf8ClearText[i*2]
		b2 := utf8ClearText[i*2+1]
		uint16Elem := (uint16(b1) << 8) | uint16(b2)
		uint16ClearTexts = append(uint16ClearTexts, uint16Elem)
	}
	if len(utf8ClearText)%2 != 0 {
		b1 := utf8ClearText[len(utf8ClearText)-1]
		b2 := 0
		uint16Elem := (uint16(b1) << 8) | uint16(b2)
		uint16ClearTexts = append(uint16ClearTexts, uint16Elem)
	}
	// 以后可以改成批量同态加密
	utf8EncText := []byte{}
	for i := 0; i < len(uint16ClearTexts); i++ {
		uint16Enc, err := FHEUint16Enc(uint16ClearTexts[i])
		if err != nil {
			return nil, err
		}
		utf8EncText = append(utf8EncText, uint16Enc...)
	}
	return utf8EncText, nil
}

func FHEUtf8StrDec(utf8EncText []byte) (string, error) {
	if len(utf8EncText)%128 != 0 {
		return "", fmt.Errorf("EncText length is not a multiple of 128")
	}
	utf8DecTextBytes := []byte{}
	for i := 0; i < len(utf8EncText)/128; i++ {
		uint16EncText := utf8EncText[i*128 : (i+1)*128]
		uint16ClearText, err := FHEUint16Dec(uint16EncText)
		if err != nil {
			return "", err
		}
		b2 := byte(uint16ClearText & 0xFF)
		b1 := byte(uint16ClearText >> 8)
		utf8DecTextBytes = append(utf8DecTextBytes, b1, b2)
	}
	var str = string(utf8DecTextBytes)
	if len(str) > 0 && int(str[len(str)-1]) == 0 {
		str = str[:len(str)-1]
	}
	return str, nil
}

var fheX0 *big.Int

func GetX0(secretKeyBasePath string) *big.Int {
	fhePubKey, err := interfaceAPI.FileReadKeyByte(secretKeyBasePath + "./05_NFHE_DECD_FHEKey.txt")
	if err != nil {
		panic(err)
	}
	// FheX0 is 1024 bit, 128 Byte.
	fheX0 = new(big.Int)
	n := len(fhePubKey)
	fheX0Bytes := fhePubKey[n-FheByteLen:]
	fheX0.SetBytes(fheX0Bytes) // fhePubKey从文件中读出，本就是大端
	return fheX0
}

const FheByteLen = 128 // 一个同态密文有128字节

func FheUnitSub(enc1, enc2 []byte) []byte {
	// enc1和enc2从加密卡中输出，是小端
	// bit.Int.SetBytes(), 参数是大端，需要先转化一下
	// 先将密文由小端转化成大端，进行大数计算
	enc1Reversed := reverseBytes(enc1) //小端转大端
	enc2Reversed := reverseBytes(enc2) //小端转大端
	enc1Num := new(big.Int)
	enc2Num := new(big.Int)
	encDiff := new(big.Int)
	enc1Num.SetBytes(enc1Reversed)
	enc2Num.SetBytes(enc2Reversed)
	if enc1Num.Cmp(enc2Num) == -1 { // enc1 < enc2
		enc1Num.Add(enc1Num, fheX0)
	}
	encDiff.Sub(enc1Num, enc2Num)
	// 大数计算完成后，由大端转化成小端，后面可以拿小端进行解密
	res := encDiff.Bytes()
	res = reverseBytes(res) // 大端转小端
	// 高位补0，对于小端来说，就是尾部补零
	curLen := len(res)
	for i := curLen; i < FheByteLen; i++ {
		res = append(res, 0)
	}
	return res
}

func reverseBytes(bin []byte) []byte {
	// 不修改原数组
	binReversed := make([]byte, len(bin))
	for i := 0; i < len(bin); i++ {
		binReversed[len(bin)-i-1] = bin[i]
	}
	return binReversed
}

func FheMultiSub(multiEnc1, multiEnc2 []byte) ([]byte, error) {
	n1 := len(multiEnc1)
	n2 := len(multiEnc2)
	if n1 != n2 {
		return nil, fmt.Errorf("multiEnc1 and multiEnc2 lengths not equal")
	}
	if n1 == 0 {
		return nil, fmt.Errorf("multiEnc1 and multiEnc2 length is 0")
	}
	if n1%FheByteLen != 0 {
		return nil, fmt.Errorf("multiEnc1 and MultiEnc2 lengths not multiple of 128")
	}
	n := n1
	res := make([]byte, 0)
	for i := 0; i < n/FheByteLen; i++ {
		res = append(res, FheUnitSub(multiEnc1[i*FheByteLen:(i+1)*FheByteLen], multiEnc2[i*FheByteLen:(i+1)*FheByteLen])...)
	}
	return res, nil
}

func FheMultiDiffDec(multiEnc []byte) (bool, error) {
	n := len(multiEnc)
	if n == 0 {
		return false, fmt.Errorf("diff length is 0")
	}
	if n%FheByteLen != 0 {
		return false, fmt.Errorf("diff length is not a multiple of %d", FheByteLen)
	}
	for i := 0; i < n/FheByteLen; i++ {
		dec, err := FHEUint16Dec(multiEnc[i*FheByteLen : (i+1)*FheByteLen])
		if err != nil {
			return false, err
		}
		if dec != 0 {
			return false, nil
		}
	}
	return true, nil
}
